Expose exact LLM token usage for agentic queries - #2580
Open
mahikaw wants to merge 2 commits into
Open
Conversation
…vider usage through Python, CLI, service, and MCP paths so callers can report inner ReAct tokens without approximation.
Contributor
Greptile SummaryThis PR collects provider-reported token usage from the ReAct and selection stages and exposes normalized totals through Python, CLI, and service responses.
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/usage.py | Adds normalized input, output, and total token accounting while preserving provider stage data. |
| nemo_retriever/src/nemo_retriever/query/agentic.py | Collects query-scoped usage from both agent operators and maps graph-position identifiers back to caller query identifiers. |
| nemo_retriever/src/nemo_retriever/query/workflow.py | Adds a metadata-returning agentic workflow while preserving the existing hits-only function. |
| nemo_retriever/src/nemo_retriever/service/query_schema.py | Defines additive typed usage fields for agentic service responses. |
| nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/atif.py | Builds bounded ATIF trajectories and persists them through a best-effort atomic file replacement. |
Sequence Diagram
sequenceDiagram
participant Caller
participant Workflow
participant ReAct
participant Selection
participant Provider
Caller->>Workflow: Agentic query
Workflow->>ReAct: Run query
ReAct->>Provider: LLM calls
Provider-->>ReAct: Results and token usage
Workflow->>Selection: Rank fused candidates
Selection->>Provider: Selection call
Provider-->>Selection: Result and token usage
Workflow->>ReAct: pop_query_usage(query_id)
Workflow->>Selection: pop_query_usage(query_id)
Workflow-->>Caller: Hits and normalized usage
Reviews (2): Last reviewed commit: "Add ATIF tracing and token usage reporti..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+30
to
+33
| @dataclass(frozen=True) | ||
| class AgenticQueryDocumentsResult: | ||
| hits: list[dict[str, Any]] | ||
| usage: dict[str, Any] |
Contributor
There was a problem hiding this comment.
The new public result classes and usage-returning methods omit field, parameter, return, and error documentation, leaving API consumers without a defined usage shape or ownership contract. The same incomplete pattern also appears on AgenticRetrieveResult, retrieve_with_usage(), and both operator pop_query_usage() methods.
Rule Used: Every public class and function in nemo_retriever ... (source)
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/query/workflow.py
Line: 30-33
Comment:
**Incomplete public usage contracts**
The new public result classes and usage-returning methods omit field, parameter, return, and error documentation, leaving API consumers without a defined usage shape or ownership contract. The same incomplete pattern also appears on `AgenticRetrieveResult`, `retrieve_with_usage()`, and both operator `pop_query_usage()` methods.
**Rule Used:** Every public class and function in nemo_retriever ... ([source](https://github.com/nvidia/nemo-retriever/blob/a48f2b12cccbcdec8ad6465b03821d6f4cc9fc80/nemo_retriever/.greptile/config.json))
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
retriever query --agentic --include-usageWhy
Consumers need exact inner-agent token accounting. Parsing logs or estimating from text does not provide reliable provider-reported usage.
Implementation
pop_query_usage()callsfinallyblocks so backend state is cleared safelyusage: nullwhen the provider does not report usageChecklist